# coding=utf-8
import os
import cv2
import numpy as np
from PIL import Image

def binaryzation(data):
    row = data.shape[1]
    col = data.shape[2]
    ret = np.empty(row * col)
    for i in range(row):
        for j in range(col):
            ret[i * col + j] = 0
            # print(data[0][i][j])
            if data[0][i][j] > 80:
                ret[i * col + j] = 1
    return ret


def load_data(data_path):
    files = os.listdir(data_path)
    file_num = len(files)
    # print(file_num)

    selected_file_num = file_num
    selected_files = []
    for i in range(selected_file_num):
        selected_files.append(files[i])

    img_mat = np.empty((selected_file_num, 1, 28, 28), dtype="float32")

    data = np.empty((selected_file_num, 28 * 28), dtype="float32")
    label = np.empty(selected_file_num, dtype="uint8")

    #print("loading data...")
    for i in range(selected_file_num):
        # print(i, "/", selected_file_num, "\r")
        file_name = selected_files[i]
        file_path = os.path.join(data_path, file_name)
        #print file_path
        img_mat[i] = Image.open(file_path)
        data[i] = binaryzation(img_mat[i])
        label[i] = int(file_name.split('.')[0])
    #print("")

    return data, label


def KNN(test_vec, train_data, train_label, k):
    train_data_size = train_data.shape[0]
    dif_mat = np.tile(test_vec, (train_data_size, 1)) - train_data
    sqr_dif_mat = dif_mat ** 2
    sqr_dis = sqr_dif_mat.sum(axis=1)

    sorted_idx = sqr_dis.argsort()

    class_cnt = {}
    maxx = 0
    best_class = 0
    for i in range(k):
        tmp_class = train_label[sorted_idx[i]]
        tmp_cnt = class_cnt.get(tmp_class, 0) + 1
        class_cnt[tmp_class] = tmp_cnt
        if (tmp_cnt > maxx):
            maxx = tmp_cnt
            best_class = tmp_class
    return best_class


def giveresult(filename):
    train_data, train_label = load_data("data")

    img_mat = np.empty((5, 1, 28, 28), dtype="float32")

    data = np.empty((5, 28 * 28), dtype="float32")
    
    img_mat[0] = Image.open('output/' + filename)
    test_data = binaryzation(img_mat[0])

    best_class = KNN(test_data, train_data, train_label, 1)

    print(best_class)

    return best_class
